import os
from flask import Flask, request, jsonify, send_file, render_template, session
from werkzeug.utils import secure_filename
from docx import Document
import uuid
from datetime import datetime
import zipfile
import io
import PyPDF2
import re
import json
from io import BytesIO
import multiprocessing
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
import traceback
import logging

# 配置日志
logging.basicConfig(level=logging.INFO, 
                    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
                    filename='app.log',
                    filemode='a')
logger = logging.getLogger(__name__)

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = 'uploads'
app.config['OUTPUT_FOLDER'] = 'outputs'
app.config['MAPPING_FILE'] = 'filename_mapping.json'
app.config['MAX_CONTENT_LENGTH'] = 500 * 1024 * 1024  # 500MB限制
app.config['SEND_FILE_MAX_AGE_DEFAULT'] = 0  # 禁用缓存
app.secret_key = 'markdown_converter_secret_key'  # 添加密钥，用于session

# 安全地确定处理器数量
try:
    app.config['MAX_WORKERS'] = min(4, multiprocessing.cpu_count())  # 限制最大进程数为4
except:
    app.config['MAX_WORKERS'] = 2  # 默认值

# 标志位，用于控制是否使用多进程
USE_MULTIPROCESSING = True  # 可以根据需要关闭

# 确保上传和输出目录存在
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
os.makedirs(app.config['OUTPUT_FOLDER'], exist_ok=True)

# 文件名映射
filename_mapping = {}

# 如果映射文件存在，加载它
if os.path.exists(app.config['MAPPING_FILE']):
    try:
        with open(app.config['MAPPING_FILE'], 'r', encoding='utf-8') as f:
            filename_mapping = json.load(f)
    except Exception as e:
        logger.error(f"加载映射文件失败: {str(e)}")

def save_mapping():
    """保存文件名映射到文件"""
    try:
        with open(app.config['MAPPING_FILE'], 'w', encoding='utf-8') as f:
            json.dump(filename_mapping, f, ensure_ascii=False, indent=2)
    except Exception as e:
        logger.error(f"保存映射文件失败: {str(e)}")

def convert_docx_to_markdown(docx_path, output_path):
    """将Word文档转换为Markdown格式"""
    try:
        doc = Document(docx_path)
        markdown_text = []
        
        # 处理段落
        for para in doc.paragraphs:
            if para.style.name.startswith('Heading'):
                level = int(para.style.name[-1])
                markdown_text.append('#' * level + ' ' + para.text + '\n')
            else:
                # 处理文本样式
                text = para.text
                runs = para.runs
                if runs:
                    formatted_text = []
                    for run in runs:
                        if run.bold:
                            formatted_text.append(f'**{run.text}**')
                        elif run.italic:
                            formatted_text.append(f'*{run.text}*')
                        else:
                            formatted_text.append(run.text)
                    text = ''.join(formatted_text)
                markdown_text.append(text + '\n')
        
        # 处理表格
        for table in doc.tables:
            # 添加表头
            header_parts = ['|']
            separator_parts = ['|']
            for cell in table.rows[0].cells:
                header_parts.append(cell.text.strip() + '|')
                separator_parts.append('---|')
            
            markdown_text.append(''.join(header_parts))
            markdown_text.append(''.join(separator_parts))
            
            # 添加数据行
            for row in table.rows[1:]:
                data_parts = ['|']
                for cell in row.cells:
                    data_parts.append(cell.text.strip() + '|')
                markdown_text.append(''.join(data_parts))
            
            markdown_text.append('\n')

        # 写入文件
        with open(output_path, 'w', encoding='utf-8') as f:
            f.write('\n'.join(markdown_text))

        return True
    except Exception as e:
        logger.error(f"Word转换失败: {str(e)}\n{traceback.format_exc()}")
        with open(output_path, 'w', encoding='utf-8') as f:
            f.write(f"# Word转换错误\n\n无法转换Word文件: {str(e)}")
        return str(e)

def convert_pdf_to_markdown_single_process(pdf_path, output_path):
    """单进程转换PDF文档为Markdown格式"""
    try:
        with open(pdf_path, 'rb') as file:
            reader = PyPDF2.PdfReader(file)
            markdown_text = []
            
            # 按顺序处理页面
            for page_num, page in enumerate(reader.pages):
                text = page.extract_text()
                
                # 简单处理章节标题
                lines = text.split('\n')
                for line in lines:
                    line = line.strip()
                    if not line:
                        continue
                        
                    # 尝试识别标题
                    if len(line) < 100 and re.match(r'^[A-Z0-9]', line) and '.' in line[:15]:
                        # 可能是标题
                        markdown_text.append(f'## {line}\n')
                    else:
                        markdown_text.append(f'{line}\n')
                
                markdown_text.append(f'\n--- 第 {page_num + 1} 页 ---\n')
                
        # 写入文件
        with open(output_path, 'w', encoding='utf-8') as f:
            f.write('\n'.join(markdown_text))
        
        return True
    except Exception as e:
        logger.error(f"单进程PDF转换失败: {str(e)}\n{traceback.format_exc()}")
        with open(output_path, 'w', encoding='utf-8') as f:
            f.write(f"# PDF转换错误\n\n无法转换PDF文件: {str(e)}")
        return str(e)

def process_pdf_page(args):
    """处理单个PDF页面的函数，避免使用闭包或实例方法"""
    page, page_num = args
    try:
        text = page.extract_text()
        markdown_lines = []
        
        # 简单处理章节标题
        lines = text.split('\n')
        for line in lines:
            line = line.strip()
            if not line:
                continue
                
            # 尝试识别标题
            if len(line) < 100 and re.match(r'^[A-Z0-9]', line) and '.' in line[:15]:
                # 可能是标题
                markdown_lines.append(f'## {line}\n')
            else:
                markdown_lines.append(f'{line}\n')
        
        markdown_lines.append(f'\n--- 第 {page_num + 1} 页 ---\n')
        return '\n'.join(markdown_lines)
    except Exception as e:
        logger.error(f"处理PDF页面 {page_num} 失败: {str(e)}")
        return f"# 页面 {page_num + 1} 处理错误\n\n无法处理此页: {str(e)}\n"

def convert_pdf_to_markdown(pdf_path, output_path):
    """多线程转换PDF文档为Markdown格式"""
    try:
        # 如果禁用了多进程，使用单进程版本
        if not USE_MULTIPROCESSING:
            return convert_pdf_to_markdown_single_process(pdf_path, output_path)
            
        with open(pdf_path, 'rb') as file:
            reader = PyPDF2.PdfReader(file)
            pages = reader.pages
            
            # 使用线程池而不是进程池，避免序列化问题
            with ThreadPoolExecutor(max_workers=app.config['MAX_WORKERS']) as executor:
                # 创建任务列表
                page_args = [(page, i) for i, page in enumerate(pages)]
                results = list(executor.map(process_pdf_page, page_args))
                
        # 写入文件
        with open(output_path, 'w', encoding='utf-8') as f:
            f.write(''.join(results))
        
        return True
    except Exception as e:
        logger.error(f"多线程PDF转换失败: {str(e)}\n{traceback.format_exc()}")
        # 尝试退回到单线程模式
        try:
            logger.info("尝试使用单线程模式转换PDF")
            return convert_pdf_to_markdown_single_process(pdf_path, output_path)
        except Exception as inner_e:
            logger.error(f"单线程模式也失败: {str(inner_e)}")
            with open(output_path, 'w', encoding='utf-8') as f:
                f.write(f"# PDF转换错误\n\n无法转换PDF文件: {str(e)}\n\n再次尝试也失败: {str(inner_e)}")
            return str(e)

def convert_file(file_data):
    """处理单个文件转换的函数"""
    file, temp_path, output_path, file_ext = file_data
    
    try:
        # 保存上传的文件
        file.save(temp_path)
        
        # 转换文件
        if file_ext == '.docx':
            result = convert_docx_to_markdown(temp_path, output_path)
        else:  # PDF
            result = convert_pdf_to_markdown(temp_path, output_path)
        
        # 删除原始文件
        try:
            os.remove(temp_path)
        except Exception as e:
            logger.warning(f"无法删除临时文件 {temp_path}: {str(e)}")
        
        return result
    except Exception as e:
        logger.error(f"文件转换失败 {file.filename}: {str(e)}\n{traceback.format_exc()}")
        with open(output_path, 'w', encoding='utf-8') as f:
            f.write(f"# 转换错误\n\n文件转换失败: {str(e)}")
        return str(e)

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/upload', methods=['POST'])
def upload_files():
    if 'files[]' not in request.files:
        return jsonify({'error': '没有文件被上传'}), 400
    
    files = request.files.getlist('files[]')
    results = []
    converted_files = []  # 用于保存本次转换的文件列表
    conversion_tasks = []  # 存储转换任务
    
    for file in files:
        if file.filename == '':
            continue
            
        file_ext = os.path.splitext(file.filename)[1].lower()
            
        if file_ext not in ['.docx', '.pdf']:
            results.append({
                'filename': file.filename,
                'status': 'error',
                'message': '仅支持.docx和.pdf文件'
            })
            continue
            
        try:
            # 生成唯一的文件名，但保留原始文件名
            original_name = os.path.splitext(file.filename)[0]
            unique_id = uuid.uuid4().hex[:8]
            
            # 创建安全的文件名
            safe_original = secure_filename(original_name)
            system_filename = f"{safe_original}_{unique_id}.md"
            temp_filename = f"{safe_original}_{unique_id}{file_ext}"
            
            temp_path = os.path.join(app.config['UPLOAD_FOLDER'], temp_filename)
            output_path = os.path.join(app.config['OUTPUT_FOLDER'], system_filename)
            
            # 添加到转换任务列表
            conversion_tasks.append((file, temp_path, output_path, file_ext))
            
            # 保存文件名映射
            filename_mapping[system_filename] = {
                'original_filename': file.filename,
                'original_name': original_name
            }
            
            # 添加到转换文件列表
            converted_files.append(system_filename)
            
            results.append({
                'filename': file.filename,
                'original_name': original_name,
                'status': 'pending',
                'md_filename': system_filename
            })
            
        except Exception as e:
            logger.error(f"准备转换文件时出错 {file.filename}: {str(e)}\n{traceback.format_exc()}")
            results.append({
                'filename': file.filename,
                'status': 'error',
                'message': str(e)
            })
    
    # 保存文件名映射
    save_mapping()
    
    # 处理转换任务
    try:
        if conversion_tasks:
            # 始终使用顺序处理模式，更可靠
            for i, task in enumerate(conversion_tasks):
                try:
                    result = convert_file(task)
                    if result is not True:
                        results[i]['status'] = 'error'
                        results[i]['message'] = str(result)
                    else:
                        results[i]['status'] = 'success'
                except Exception as e:
                    logger.error(f"处理文件时出错: {str(e)}\n{traceback.format_exc()}")
                    results[i]['status'] = 'error'
                    results[i]['message'] = str(e)
    except Exception as e:
        logger.error(f"批量处理文件时出错: {str(e)}\n{traceback.format_exc()}")
        return jsonify({'error': f'处理文件时发生错误: {str(e)}'}), 500
    
    # 保存本次转换的文件列表到session
    session['converted_files'] = converted_files
    
    return jsonify(results)

@app.route('/download/<filename>')
def download_file(filename):
    try:
        file_path = os.path.join(app.config['OUTPUT_FOLDER'], filename)
        
        # 确保文件存在
        if not os.path.exists(file_path):
            return jsonify({'error': '文件不存在'}), 404
        
        # 从映射中获取原始文件名
        if filename in filename_mapping:
            # 获取原始完整文件名并替换扩展名为.md
            original_filename = filename_mapping[filename]['original_filename']
            base_name = os.path.splitext(original_filename)[0]
            download_name = f"{base_name}.md"
        else:
            # 如果没有找到映射，使用安全的名称
            download_name = filename.split('_')[0] + '.md'
            
        return send_file(
            file_path,
            as_attachment=True,
            download_name=download_name
        )
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/download-all')
def download_all():
    try:
        # 获取当前时间戳作为ZIP文件名的一部分
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        zip_filename = f"markdown_files_{timestamp}.zip"
        memory_file = BytesIO()
        
        with zipfile.ZipFile(memory_file, 'w') as zf:
            # 获取会话中的最近转换文件列表
            session_files = session.get('converted_files', [])
            
            if not session_files:
                # 如果会话中没有文件，则默认包含所有文件
                for filename in os.listdir(app.config['OUTPUT_FOLDER']):
                    file_path = os.path.join(app.config['OUTPUT_FOLDER'], filename)
                    if os.path.isfile(file_path):
                        with open(file_path, 'rb') as f:
                            # 获取原始文件名用于ZIP中的文件
                            if filename in filename_mapping:
                                original_filename = filename_mapping[filename]['original_filename']
                                base_name = os.path.splitext(original_filename)[0]
                                markdown_filename = f"{base_name}.md"
                            else:
                                # 如果没有找到映射，使用当前文件名
                                markdown_filename = filename
                            
                            # 将文件添加到ZIP中
                            zf.writestr(markdown_filename, f.read())
            else:
                # 只包含会话中的最近转换文件
                for system_filename in session_files:
                    file_path = os.path.join(app.config['OUTPUT_FOLDER'], system_filename)
                    if os.path.isfile(file_path):
                        with open(file_path, 'rb') as f:
                            # 获取原始文件名用于ZIP中的文件
                            if system_filename in filename_mapping:
                                original_filename = filename_mapping[system_filename]['original_filename']
                                base_name = os.path.splitext(original_filename)[0]
                                markdown_filename = f"{base_name}.md"
                            else:
                                # 如果没有找到映射，使用当前文件名
                                markdown_filename = system_filename
                            
                            # 将文件添加到ZIP中
                            zf.writestr(markdown_filename, f.read())
        
        memory_file.seek(0)
        return send_file(
            memory_file, 
            mimetype='application/zip', 
            as_attachment=True, 
            download_name=zip_filename
        )
    except Exception as e:
        return jsonify({'error': str(e)}), 500

if __name__ == '__main__':
    app.run(debug=True, host='127.0.0.1', port=5000, threaded=True) 